import { isString } from 'lodash-es'
import { dirname, resolve } from 'path'
import { decode } from 'sourcemap-codec'
import type {
  AcornNode,
  Awaitable,
  DecodedSourceMapOrMissing,
  ExistingDecodedSourceMap,
  ExistingRawSourceMap,
  SourceMapInput,
  SourceMapSegment,
  TransformResult
} from './types'
import { doThen, isAst } from './utils'

export const decodedSourcemap = (map: SourceMapInput | ExistingDecodedSourceMap): ExistingDecodedSourceMap | null => {
  if (!map) return null

  if (isString(map)) {
    map = JSON.parse(map) as ExistingRawSourceMap
  }
  if (map.mappings === '') {
    return {
      mappings: [],
      names: [],
      sources: [],
      version: 3
    }
  }

  const mappings = typeof map.mappings === 'string' ? decode(map.mappings) : map.mappings

  return { ...(map as ExistingRawSourceMap | ExistingDecodedSourceMap), mappings }
}

class Source {
  readonly content: string
  readonly filename: string
  isOriginal = true

  constructor(filename: string, content: string) {
    this.filename = filename
    this.content = content
  }

  traceSegment(line: number, column: number, name: string): SourceMapSegmentObject {
    return { column, line, name, source: this }
  }
}

interface SourceMapSegmentObject {
  column: number
  line: number
  name: string
  source: Source
}

class Link {
  readonly mappings: readonly SourceMapSegment[][]
  readonly names: readonly string[]
  readonly sources: (Source | Link)[]

  constructor(map: { mappings: readonly SourceMapSegment[][]; names: readonly string[] }, sources: (Source | Link)[]) {
    this.sources = sources
    this.names = map.names
    this.mappings = map.mappings
  }

  traceMappings() {
    const sources: string[] = []
    const sourceIndexMap = new Map<string, number>()
    const sourcesContent: string[] = []
    const names: string[] = []
    const nameIndexMap = new Map<string, number>()

    const mappings = []

    for (const line of this.mappings) {
      const tracedLine: SourceMapSegment[] = []

      for (const segment of line) {
        if (segment.length === 1) continue
        const source = this.sources[segment[1]]
        if (!source) continue

        const traced = source.traceSegment(segment[2], segment[3], segment.length === 5 ? this.names[segment[4]] : '')

        if (traced) {
          const {
            column,
            line,
            name,
            source: { content, filename }
          } = traced
          let sourceIndex = sourceIndexMap.get(filename)
          if (sourceIndex === undefined) {
            sourceIndex = sources.length
            sources.push(filename)
            sourceIndexMap.set(filename, sourceIndex)
            sourcesContent[sourceIndex] = content
          } else if (sourcesContent[sourceIndex] == null) {
            sourcesContent[sourceIndex] = content
          } else if (content != null && sourcesContent[sourceIndex] !== content) {
            throw new Error(`Multiple conflicting contents for sourcemap source ${filename}`)
          }

          const tracedSegment: SourceMapSegment = [segment[0], sourceIndex, line, column]

          if (name) {
            let nameIndex = nameIndexMap.get(name)
            if (nameIndex === undefined) {
              nameIndex = names.length
              names.push(name)
              nameIndexMap.set(name, nameIndex)
            }

            ;(tracedSegment as SourceMapSegment)[4] = nameIndex
          }

          tracedLine.push(tracedSegment)
        }
      }

      mappings.push(tracedLine)
    }

    return { mappings, names, sources, sourcesContent }
  }

  traceSegment(line: number, column: number, name: string): SourceMapSegmentObject | null {
    const segments = this.mappings[line]
    if (!segments) return null

    // binary search through segments for the given column
    let searchStart = 0
    let searchEnd = segments.length - 1

    while (searchStart <= searchEnd) {
      const m = (searchStart + searchEnd) >> 1
      const segment = segments[m]

      // If a sourcemap does not have sufficient resolution to contain a
      // necessary mapping, e.g. because it only contains line information, we
      // use the best approximation we could find
      if (segment[0] === column || searchStart === searchEnd) {
        if (segment.length == 1) return null
        const source = this.sources[segment[1]]
        if (!source) return null

        return source.traceSegment(segment[2], segment[3], segment.length === 5 ? this.names[segment[4]] : name)
      }
      if (segment[0] > column) {
        searchEnd = m - 1
      } else {
        searchStart = m + 1
      }
    }

    return null
  }
}

function linkMap(source: Source | Link, map: DecodedSourceMapOrMissing): Link {
  if (map.mappings) {
    return new Link(map, [source])
  }

  return new Link(
    {
      mappings: [],
      names: []
    },
    [source]
  )
}

export function getCollapsedSourcemap(
  id: string,
  originalCode: string,
  originalSourcemap: ExistingDecodedSourceMap | null,
  sourcemapChain: readonly DecodedSourceMapOrMissing[]
): Source | Link {
  let source: Source | Link

  if (!originalSourcemap) {
    source = new Source(id, originalCode)
  } else {
    const sources = originalSourcemap.sources
    const sourcesContent = originalSourcemap.sourcesContent || []
    const directory = dirname(id) || '.'
    const sourceRoot = originalSourcemap.sourceRoot || '.'

    const baseSources = sources.filter(Boolean).map(
      (source, i) => new Source(resolve(directory, sourceRoot, source), sourcesContent[i])
    )
    source = new Link(originalSourcemap, baseSources)
  }
  return sourcemapChain.reduce(linkMap, source)
}

export function collapseSourcemap(
  id: string,
  originalCode: string,
  originalSourcemap: ExistingDecodedSourceMap | null,
  sourcemapChain: readonly DecodedSourceMapOrMissing[]
): ExistingDecodedSourceMap | null {
  if (!sourcemapChain.length) {
    return originalSourcemap
  }

  const source = getCollapsedSourcemap(id, originalCode, originalSourcemap, sourcemapChain) as Link
  const map = source.traceMappings()
  return { version: 3, ...map }
}

export const doTransform = (
  id: string,
  code: string,
  map: ExistingDecodedSourceMap | false | null | undefined,
  ast: AcornNode | undefined,
  transforms: Iterable<(code: string, id: string, ast?: AcornNode) => Awaitable<TransformResult>>
) => {
  const sourcemapChain: DecodedSourceMapOrMissing[] | undefined =
    map === false ? undefined : [decodedSourcemap(map) || { missing: true }]
  let awaitableCodeAndAst: Awaitable<[code: string, ast?: AcornNode]> = [code, ast]

  for (const transform of transforms) {
    awaitableCodeAndAst = doThen(awaitableCodeAndAst, previous => {
      const [previousCode, previousAst] = previous
      return doThen(transform(previousCode, id, previousAst), result => {
        let code: string
        let map: string | ExistingRawSourceMap | { mappings: '' } | null | undefined
        let ast: AcornNode | undefined

        if (result == null) {
          return previous
        } else if (isString(result)) {
          code = String(result)
        } else if ((result as any)?.code != null) {
          ;({ code, map, ast } = result)

          code = String(code)

          if (!isAst(ast)) {
            ast = undefined
          }
        } else {
          return previous
        }

        // strict null check allows 'null' maps to not be pushed to the chain,
        // while 'undefined' gets the missing map warning
        if (map != null) {
          sourcemapChain?.push(decodedSourcemap(map as any) || { missing: true })
        }

        return [code, ast]
      })
    })
  }

  return doThen(awaitableCodeAndAst, ([transformedCode, transformAst]) => {
    return {
      code: transformedCode,
      map: sourcemapChain ? getCollapsedSourcemap(id, code, map as any, sourcemapChain) : undefined,
      ast: transformAst
    }
  })
}
